#include <iostream>
#include <vector>
#include <cassert>
#include "MinHeap.h"

using namespace std;


template<typename Graph, typename Weight>
class LazyPrimMST{

private:
    Graph &G;
    MinHeap<Edge<Weight>> pq;
    bool *marked;					// 点是否被标记
    vector<Edge<Weight>> mst;		// 存储 n-1个边到MST向量中去
    Weight mstWeight;
	
	/**
	 * 不断将蓝色节点转化为红色节点
	 * 遍历相应节点边,将新的横切边加入到最小堆中
	 * */
    void visit(int v){

        assert( !marked[v] );
        marked[v] = true;				// 访问后将其变为红色

        typename Graph::adjIterator adj(G,v);
        for( Edge<Weight>* e = adj.begin() ; !adj.end() ; e = adj.next() )
			// 如果v对应节点没有被marked，将两点对应边加入到最小堆中
            if( !marked[e->other(v)] )
                pq.insert(*e);
    }

public:
    LazyPrimMST(Graph &graph):G(graph), pq(MinHeap<Edge<Weight>>(graph.E())){

        marked = new bool[G.V()];
        for( int i = 0 ; i < G.V() ; i ++ )
            marked[i] = false;
        mst.clear();

        // Lazy Prim	O(ElogE)
        visit(0);
        while( !pq.isEmpty() ){						// 优先队列不为空
            Edge<Weight> e = pq.extractMin();		// 取出优先队列中最小边 
			// (两个点同时被标记)如果在切分的同侧,不是切分线
            if( marked[e.v()] == marked[e.w()] )   	
                continue;

			// 否则从优先队列中取出最小边,同时也是横切边,将其push仅mst向量中
            mst.push_back( e );			
			// 此时e中一个点红色一个点蓝色,
            if( !marked[e.v()] )	// e.v 未被标记,是蓝色端点
                visit( e.v() );
            else					// e.w 是蓝色的端点
                visit( e.w() );
        }

        mstWeight = mst[0].wt();	// 计算mst最小生成树权值
        for( int i = 1 ; i < mst.size() ; i ++ )
            mstWeight += mst[i].wt();
    }

    ~LazyPrimMST(){
        delete[] marked;
    }
	// 返回mst向量,最小生成树
    vector<Edge<Weight>> mstEdges(){
        return mst;
    };

	// 返回最小生成树的最小权值
    Weight result(){
        return mstWeight;
    };
};
